import cv2
from ultralytics import YOLO
from typing import Optional, Tuple, List
import numpy as np
from pathlib import Path
from .models import BBox, Detection, FrameDetections


class VideoObjectDetector:
    def __init__(self, model_path: str = 'yolov8n.pt'):
        self.model_path = model_path
        self.model = YOLO(model_path)
        self.classes = self.model.names

    def detect_video(
            self,
            video_path: str,
            output_path: Optional[str] = None,
            conf_threshold: float = 0.5,
            show: bool = False,
            save: bool = True,
            fps: Optional[float] = None
    ) -> Tuple[bool, str, List[FrameDetections]]:
        """
        处理视频文件，逐帧检测目标。

        参数:
        video_path：输入视频文件的路径。
        output_path：输出视频的保存路径。如果None，将保存到‘output_{input_path}’。
        conf_threshold：检测的置信阈值。
        show：检测过程中是否显示视频。
        save：是否保存输出视频。
        fps：输出视频的每秒帧数。如果None，则使用输入视频的FPS。

        返回:
        元组包含：
        - success：布尔值，表示处理是否成功
        - message：包含状态消息或错误的字符串
        -
        """
        try:
            # Open the video file
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                return False, f"Error: Could not open video file {video_path}", []

            original_fps = cap.get(cv2.CAP_PROP_FPS)
            frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

            output_fps = fps if fps is not None else original_fps

            if output_path is None:
                input_path = Path(video_path)
                output_path = str(input_path.parent / f"output_{input_path.name}")

            if save:
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                out = cv2.VideoWriter(output_path, fourcc, output_fps, (frame_width, frame_height))

            frame_count = 0
            video_detections = []

            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break

                # 当前帧时间戳（秒）
                timestamp = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0

                # 检测
                results = self.model(frame, conf=conf_threshold)
                annotated_frame = results[0].plot()

                # 提取检测结果
                detections = []
                boxes = results[0].boxes
                if boxes is not None and boxes.data is not None:
                    for det in boxes.data.tolist():
                        x1, y1, x2, y2, conf, class_id = det
                        detection = Detection(
                            class_id=int(class_id),
                            confidence=float(conf),
                            bbox=BBox(float(x1), float(y1), float(x2), float(y2))
                        )
                        detections.append(detection)

                # 封装当前帧的检测结果
                frame_detections = FrameDetections(timestamp=timestamp, detections=detections)
                video_detections.append(frame_detections)

                # 显示和保存
                if show:
                    cv2.imshow('YOLOv8 Detection', annotated_frame)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                if save:
                    out.write(annotated_frame)

                frame_count += 1
                print(f"Processed frame {frame_count}/{total_frames}", end='\r')

            cap.release()
            if save:
                out.release()
            if show:
                cv2.destroyAllWindows()

            return True, f"Successfully processed video. Output saved to {output_path}", video_detections

        except Exception as e:
            return False, f"Error processing video: {str(e)}", []

    def get_class_names(self) -> list:
        """Returns the list of class names that the model can detect."""
        return list(self.classes.values())